# https://blog.csdn.net/weixin_39228381/article/details/108511882   原文
import numpy as np
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
 
def func(x, y):
    return x * x + 10 * y * y
 
def paint_loss_func():
    x = np.linspace(-50, 50, 100) #x的绘制范围是-50到50，从改区间均匀取100个数
    y = np.linspace(-50, 50, 100) #y的绘制范围是-50到50，从改区间均匀取100个数
    
    X, Y = np.meshgrid(x, y)
    Z = func(X, Y)
 
    fig = plt.figure()#figsize=(10, 10))
    ax = Axes3D(fig)
    plt.xlabel('x')
    plt.ylabel('y')
 
    ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap='rainbow')
    plt.show()
 
paint_loss_func()

def grad(x, y): #根据上述代码，可知x和y的梯度分别为2x和20y
    return 2 * x, 20 * y
 
def train_SGD():
    cur_x = 40
    cur_y = 20
    lr = 0.096
    track_x = [cur_x] #记录x每次的值
    track_y = [cur_y] #记录y每次的值
    for i in range(10): #作为demo，这里只训练10次
        grad_x, grad_y = grad(cur_x, cur_y) #等效于神经网络的反向传播，求取各参数的梯度
        cur_x -= lr * grad_x
        cur_y -= lr * grad_y
        track_x.append(cur_x)
        track_y.append(cur_y)
    #print(track_x)
    #print(track_y)
    return track_x, track_y
 
def train_RMSProp():
    cur_x = 40
    cur_y = 20
    lr = 3
    r_x, r_y = 0, 0 #伪代码中的r
    alpha = 0.9
    eps = 1e-06
    track_x = [cur_x]
    track_y = [cur_y]
    for i in range(10):
        grad_x, grad_y = grad(cur_x, cur_y)
        r_x = alpha * r_x + (1 - alpha) * (grad_x * grad_x)
        r_y = alpha * r_y + (1 - alpha) * (grad_y * grad_y)
        cur_x -= (grad_x / (np.sqrt(r_x) + eps)) * lr
        cur_y -= (grad_y / (np.sqrt(r_y) + eps)) * lr
        track_x.append(cur_x)
        track_y.append(cur_y)
    #print(track_x)
    #print(track_y)
    return track_x, track_y
 
def paint_tracks(track_x1, track_y1, track_x2, track_y2):
    x = np.linspace(-50, 50, 100)
    y = np.linspace(-50, 50, 100)
    
    X, Y = np.meshgrid(x, y)
    Z = func(X, Y)
 
    fig = plt.figure(figsize=(10, 10))
    ax = Axes3D(fig)
    plt.xlabel('x')
    plt.ylabel('y')
 
    #ax.plot(track_x, track_y, func(np.array(track_x), np.array(track_y)), 'r--')
    tx1, ty1 = track_x1[0], track_y1[0]
    for i in range(1, len(track_x1)):
        tx2, ty2 = track_x1[i], track_y1[i]
        tx = np.linspace(tx1, tx2, 100)
        ty = np.linspace(ty1, ty2, 100)
        ax.plot(tx, ty, func(tx, ty), 'r-')
        tx1, ty1 = tx2, ty2
    ax.scatter(track_x1, track_y1, func(np.array(track_x1), np.array(track_y1)), s=50, c='r')
 
    tx1, ty1 = track_x2[0], track_y2[0]
    for i in range(1, len(track_x2)):
        tx2, ty2 = track_x2[i], track_y2[i]
        tx = np.linspace(tx1, tx2, 100)
        ty = np.linspace(ty1, ty2, 100)
        ax.plot(tx, ty, func(tx, ty), 'b-')
        tx1, ty1 = tx2, ty2
    ax.scatter(track_x2, track_y2, func(np.array(track_x2), np.array(track_y2)), s=50, c='b')
    #ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap='rainbow')
    plt.show()
 
track_x_sgd, track_y_sgd = train_SGD()
#paint_track(track_x_sgd, track_y_sgd)
 
track_x_rms, track_y_rms = train_RMSProp()
paint_tracks(track_x_sgd, track_y_sgd, track_x_rms, track_y_rms)